from PyQt5 import QtCore
from PyQt5.QtWidgets import QMainWindow, QFileDialog, QMessageBox
from PyQt5.QtCore import pyqtSignal
from designer.main_widget import Ui_MainWindow
from designer.threads import ConvertThread
import pathlib


class Window(QMainWindow):
    output = pyqtSignal(str)
    progress = pyqtSignal(int)

    def __init__(self):
        super().__init__()
        self.ui = Ui_MainWindow()
        self.ui.setupUi(self)

        self.init_data()
        self.init_ui()
        self.init_slots()
        self.show()

    def init_data(self):
        self.convert_thread = None
        self.last_dir = './'
        self.dataset_list = []

    def init_slots(self):
        self.ui.add_button.clicked.connect(self.add_dataset)
        self.ui.remove_button.clicked.connect(self.remove_dataset)
        self.ui.convert_button.clicked.connect(self.convert)
        self.progress.connect(self.progress_slot)
        self.output.connect(self.output_slot)

    def init_ui(self):
        self.ui.progress_bar.setValue(0)
        self.ui.train_rate.setValue(0.8)
        self.ui.train_rate.setSingleStep(0.01)
        self.ui.train_rate.setMaximum(1)
        self.ui.train_rate.setMinimum(0)
        self.ui.dataset_list.setModel(QtCore.QStringListModel(self.dataset_list))

    def add_dataset(self):
        dataset = QFileDialog.getExistingDirectory(self, "选择数据集", self.last_dir)
        self.last_dir = str(pathlib.Path(dataset).parent)
        if dataset:
            self.dataset_list.append(dataset)
            self.ui.dataset_list.setModel(QtCore.QStringListModel(self.dataset_list))

    def remove_dataset(self):
        index = self.ui.dataset_list.currentIndex()
        if index.isValid():
            self.dataset_list.pop(index.row())
            self.ui.dataset_list.setModel(QtCore.QStringListModel(self.dataset_list))

    def output_slot(self, msg):
        self.ui.output.append(msg)

    def progress_slot(self, value):
        self.ui.progress_bar.setValue(value)

    def convert(self):
        train_rate = self.ui.train_rate.value()
        if not self.dataset_list:
            QMessageBox.warning(self, "警告", "请添加数据集")
            return
        output_path = QFileDialog.getExistingDirectory(self, "选择输出路径", self.last_dir)
        if output_path:
            self.convert_thread = ConvertThread(self, self.dataset_list, train_rate, output_path)
            self.convert_thread.start()
            self.disable_ui()

    def disable_ui(self):
        self.ui.add_button.setEnabled(False)
        self.ui.remove_button.setEnabled(False)
        self.ui.convert_button.setEnabled(False)
        self.ui.train_rate.setEnabled(False)
        self.ui.dataset_list.setEnabled(False)

    def enable_ui(self):
        self.ui.add_button.setEnabled(True)
        self.ui.remove_button.setEnabled(True)
        self.ui.convert_button.setEnabled(True)
        self.ui.train_rate.setEnabled(True)
        self.ui.dataset_list.setEnabled(True)
